#coding:utf-8

#从图片生成数据库
#为cnn_custom_simple生成所需的`custom_data.npz`数据文件

import numpy as np
import os
from scipy.misc import imread,imresize
import matplotlib.pyplot as plt

cwd=os.getcwd()
print('all libiaries is loaded.....')
print('current folder is [%s]'%cwd)

#config
imgsize=[64,64]
use_gray=True
data_name="custom_data"


#data path
paths = ["datas/gendataset/celebs/Arnold_Schwarzenegger"
    , "datas/gendataset/celebs/Junichiro_Koizumi"
    , "datas/gendataset/celebs/Vladimir_Putin"
    , "datas/gendataset/celebs/George_W_Bush"]
categories = ['Terminator', 'Koizumi', 'Putin', 'Bush']

print('当前图片路径：')
for i ,path in enumerate(paths):
    print("[%d/%d] %s"%(i,len(paths),path))
print("数据将会存在以下路径:\n[%s]"%(cwd+'/res/'+data_name+'.npz'))


def rgb2gray(rgb):
    if len(rgb.shape) is 3:
        return np.dot(rgb[...,:3],[0.299, 0.587, 0.114])
    else:
        return rgb


nclass     = len(paths)
valid_exts = [".jpg",".gif",".png",".tga", ".jpeg"]
imgcnt     = 0
totalImg=None
totallabel=None

for i,relpath in zip(range(nclass),paths):
    path=cwd+'/'+relpath
    flist=os.listdir(path)
    #print(flist)
    for file in flist:
        if os.path.splitext(file)[1].lower() not in valid_exts:
            #this is not  file a picture
            continue
        fullpath=os.path.join(path,file)
        #get this picture absolute path
        #print(fullpath)
        curImg=imread(fullpath)
        # read this image
        if use_gray:
            grayImg=rgb2gray(curImg)
        else:
            grayImg=curImg
        #scale
        graySmall=imresize(grayImg,[imgsize[0],imgsize[1]])/255
        #print(graySmall)
        grayVector=np.reshape(graySmall,(1,-1))
        #print(grayVector)

        #store
        currlabel=np.eye(nclass,nclass)[i:i+1,:]
        #print(currlabel)
        if imgcnt is 0:
            totalImg=grayVector
            totallabel=currlabel
        else:
            totalImg=np.concatenate((totalImg,grayVector),axis=0)
            totallabel=np.concatenate((totallabel,currlabel),axis=0)
        imgcnt+=1

print("total images is :%d"%imgcnt)

def print_shape(string, x):
    print ("SHAPE OF [%s] IS [%s]" % (string, x.shape,))

#生成 指定size的数据，同时指定数据元素的范围是imgcnt
randidx=np.random.randint(imgcnt,size=imgcnt)
#print(randidx)
trainidx   = randidx[0:int(4*imgcnt/5)]
testidx    = randidx[int(4*imgcnt/5):imgcnt]
trainimg   = totalImg[trainidx, :]
trainlabel = totallabel[trainidx, :]
testimg    = totalImg[testidx, :]
testlabel  = totallabel[testidx, :]
print_shape("totalimg", totalImg)
print_shape("totallabel", totallabel)
print_shape("trainimg", trainimg)
print_shape("trainlabel", trainlabel)
print_shape("testimg", testimg)
print_shape("testlabel", testlabel)

savepath=cwd + "/datas/gendataset/" + data_name + ".npz"
np.savez(savepath,
         trainimg=trainimg,
         trainlabel=trainlabel,
         testimg=testimg,
         testlabel=testlabel,
         imgsize=imgsize,
         use_gray=use_gray,
         categories=categories
         )
print ("SAVED TO [%s]" % (savepath))

loadpath=cwd + "/datas/gendataset/" + data_name + ".npz"
l = np.load(loadpath)
print (l.files)

# 解析数据
trainimg_loaded   = l['trainimg']
trainlabel_loaded = l['trainlabel']
testimg_loaded    = l['testimg']
testlabel_loaded  = l['testlabel']
categories_loaded = l['categories']

print ("[%d] TRAINING IMAGES" % (trainimg_loaded.shape[0]))
print ("[%d] TEST IMAGES" % (testimg_loaded.shape[0]))
print ("LOADED FROM [%s]" % (savepath))


# ## 绘制载入数据

# In[8]:

ntrain_loaded = trainimg_loaded.shape[0]
batch_size = 5;
randidx = np.random.randint(ntrain_loaded, size=batch_size)
for i in randidx:
    currimg = np.reshape(trainimg_loaded[i, :], (imgsize[0], -1))
    currlabel_onehot = trainlabel_loaded[i, :]
    currlabel = np.argmax(currlabel_onehot)
    if use_gray:
        currimg = np.reshape(trainimg[i, :], (imgsize[0], -1))
        plt.matshow(currimg, cmap=plt.get_cmap('gray'))
        plt.colorbar()
    else:
        currimg = np.reshape(trainimg[i, :], (imgsize[0], imgsize[1], 3))
        plt.imshow(currimg)
    title_string = ("[%d] CLASS-%d (%s)"
                    % (i, currlabel, categories_loaded[currlabel]))
    plt.title(title_string)
    plt.show()